# Copyright 2025 the LlamaFactory team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from typing import TYPE_CHECKING, Any, Optional, TypedDict

import torch
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoModelForSeq2SeqLM,
    AutoModelForTextToWaveform,
    AutoModelForVision2Seq,
    AutoProcessor,
    AutoTokenizer,
)
from trl import AutoModelForCausalLMWithValueHead

from ..extras import logging
from ..extras.misc import count_parameters, skip_check_imports, try_download_model_from_other_hub
from ..extras.packages import is_transformers_version_greater_than
from .adapter import init_adapter
from .model_utils.liger_kernel import apply_liger_kernel
from .model_utils.misc import register_autoclass
from .model_utils.mod import convert_pretrained_model_to_mod, load_mod_pretrained_model
from .model_utils.unsloth import load_unsloth_pretrained_model
from .model_utils.valuehead import load_valuehead_params
from .patcher import patch_config, patch_model, patch_processor, patch_tokenizer, patch_valuehead_model


if is_transformers_version_greater_than("4.46.0"):
    from transformers import AutoModelForImageTextToText


if TYPE_CHECKING:
    from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer, ProcessorMixin

    from ..hparams import FinetuningArguments, ModelArguments


logger = logging.get_logger(__name__)


class TokenizerModule(TypedDict):
    tokenizer: "PreTrainedTokenizer"
    processor: Optional["ProcessorMixin"]


# 定义函数_get_init_kwargs，用于获取加载配置/分词器/模型所需的参数字典
def _get_init_kwargs(model_args: "ModelArguments") -> dict[str, Any]:
    r"""Get arguments to load config/tokenizer/model.

    Note: including inplace operation of model_args.
    """
    # 函数说明：获取加载config/tokenizer/model的参数。注意：包括model_args的就地操作。
    # 跳过检查导入的步骤，避免不必要的依赖检查时间消耗
    # 此处假定所有必要的模块已经被正确导入
    skip_check_imports()
    # 尝试从其他模型中心下载模型（如果本地不存在），并更新model_args中的模型路径信息
    model_args.model_name_or_path = try_download_model_from_other_hub(model_args)
    
    # 返回一个包含加载模型所需初始化参数的字典，这些参数来自于model_args
    return {
        # 是否信任远程代码，默认为False。设置True允许加载自定义模型架构
        "trust_remote_code": model_args.trust_remote_code,
        # 指定缓存目录，用于存储下载的模型文件和其他资源
        "cache_dir": model_args.cache_dir,
        # Git版本号或分支名，指定加载模型的具体版本
        "revision": model_args.model_revision,
        # Hugging Face Hub的访问令牌，用于私有模型的访问权限验证
        "token": model_args.hf_hub_token,
    }


# 加载分词器
# 定义 load_tokenizer 函数，用于加载预训练分词器和可选处理器
def load_tokenizer(model_args: "ModelArguments") -> "TokenizerModule":
    r"""Load pretrained tokenizer and optionally loads processor.

    Note: including inplace operation of model_args.
    """
    # 获取初始化参数字典，如 trust_remote_code、revision 等
    init_kwargs = _get_init_kwargs(model_args)
    try:
        # 尝试根据模型路径加载分词器，设置 use_fast 和 split_special_tokens 参数
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.model_name_or_path,
            use_fast=model_args.use_fast_tokenizer,
            split_special_tokens=model_args.split_special_tokens,
            # 设置填充在右侧，适用于 decoder-only 模型
            padding_side="right",
            **init_kwargs,
        )
    except ValueError:  # try the fast one
        # 如果失败，尝试强制使用 fast 分词器
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.model_name_or_path,
            # 强制启用 fast 分词器
            use_fast=True,
            padding_side="right",
            **init_kwargs,
        )
    except Exception as e:
        # 其他异常则抛出 OSError
        raise OSError("Failed to load tokenizer.") from e

    # 对分词器进行 patch，添加特殊 token 或调整配置
    patch_tokenizer(tokenizer, model_args)
    try:
        # 尝试加载处理器（如 CLIPProcessor）
        processor = AutoProcessor.from_pretrained(model_args.model_name_or_path, **init_kwargs)
        # 对处理器进行 patch，适配当前 tokenizer 和参数
        patch_processor(processor, tokenizer, model_args)
    except Exception as e:
        # 加载或处理失败时记录日志并设为 None
        logger.info_rank0(f"Failed to load processor: {e}.")
        processor = None

    # Avoid load tokenizer, see:
    # https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/auto/processing_auto.py#L324
    # 避免加载错误类型的 Processor（如某些模型返回的是 FeatureExtractor 而非 Processor）
    if processor is not None and "Processor" not in processor.__class__.__name__:
        logger.debug("The loaded processor is not an instance of Processor. Dropping it.")
        processor = None

    # 返回包含 tokenizer 和 processor 的字典
    return {"tokenizer": tokenizer, "processor": processor}


def load_config(model_args: "ModelArguments") -> "PretrainedConfig":
    r"""Load model config."""
    # 获取加载模型配置所需的初始化参数（如 trust_remote_code、revision 等）
    init_kwargs = _get_init_kwargs(model_args)
    # 使用 AutoConfig 从预训练模型路径加载配置文件（config.json）
    return AutoConfig.from_pretrained(model_args.model_name_or_path, **init_kwargs)


def load_model(
    tokenizer: "PreTrainedTokenizer",
    model_args: "ModelArguments",
    finetuning_args: "FinetuningArguments",
    is_trainable: bool = False,
    add_valuehead: bool = False,
) -> "PreTrainedModel":
    r"""Load pretrained model."""
    # 初始化模型加载所需的参数字典，包含设备、dtype等配置
    init_kwargs = _get_init_kwargs(model_args)
    # 加载预训练模型的配置文件（config.json）
    config = load_config(model_args)
    # 对配置进行补丁处理，适配当前任务、tokenizer、训练设置等
    patch_config(config, tokenizer, model_args, init_kwargs, is_trainable)
    
    # 应用 Liger Kernel 优化，提升性能（非必要时可跳过）
    apply_liger_kernel(config, model_args, is_trainable, require_logits=(finetuning_args.stage not in ["pt", "sft"]))

    # 初始化模型变量和 lazy_load 标志位
    model = None
    lazy_load = False
    # 如果使用 unsloth 并且处于训练模式或有 adapter 路径，则延迟加载
    if model_args.use_unsloth:
        if model_args.adapter_name_or_path is not None:
            lazy_load = True
        elif is_trainable:
            # 使用 unsloth 加载预训练模型以加速训练
            model = load_unsloth_pretrained_model(config, model_args)

    # 如果未提前加载且不需要延迟加载，则按常规方式加载模型
    if model is None and not lazy_load:
        # 将 config 注入到初始化参数中
        init_kwargs["config"] = config
        # 设置预训练模型路径
        init_kwargs["pretrained_model_name_or_path"] = model_args.model_name_or_path

        # 如果启用 MoD（Mixture of Depths），则使用特定加载方式
        if model_args.mixture_of_depths == "load":
            model = load_mod_pretrained_model(**init_kwargs)
        
        else:
            # 自动选择合适的模型加载类：根据 config 类型判断是文本、图像、语音等任务
            if type(config) in AutoModelForVision2Seq._model_mapping.keys():  # image-text
                load_class = AutoModelForVision2Seq
            elif (
                is_transformers_version_greater_than("4.46.0")
                and type(config) in AutoModelForImageTextToText._model_mapping.keys()
            ):  # image-text
                load_class = AutoModelForImageTextToText
            elif type(config) in AutoModelForSeq2SeqLM._model_mapping.keys():  # audio-text
                load_class = AutoModelForSeq2SeqLM
            elif type(config) in AutoModelForTextToWaveform._model_mapping.keys():  # audio hack for qwen2_5_omni
                load_class = AutoModelForTextToWaveform
            else:
                load_class = AutoModelForCausalLM

            # 判断是否从头开始训练，决定是使用 from_config 还是 from_pretrained
            if model_args.train_from_scratch:
                model = load_class.from_config(config, trust_remote_code=model_args.trust_remote_code)
            else:
                model = load_class.from_pretrained(**init_kwargs)
                # 特殊处理 Qwen2.5 Omni 模型，仅加载 thinker 部分
                if getattr(model.config, "model_type", None) == "qwen2_5_omni":
                    model = model.thinker  # use part of Omni model

        # 如果启用 MoD（转换模式），则将现有模型转换为 MoD 架构
        if model_args.mixture_of_depths == "convert":
            model = convert_pretrained_model_to_mod(model, config, model_args)

    # 如果不是延迟加载，则对模型进行补丁处理，添加适配器或价值头支持
    if not lazy_load:
        patch_model(model, tokenizer, model_args, is_trainable, add_valuehead)
        # 注册自动加载类，确保保存/加载时能正确恢复模型结构
        register_autoclass(config, model, tokenizer)

    # 初始化适配器（如 LoRA、IA³ 等），用于参数高效微调
    model = init_adapter(config, model, model_args, finetuning_args, is_trainable)

    # 如果需要添加价值头（value head），则包装模型并加载对应权重
    if add_valuehead:
        model = AutoModelForCausalLMWithValueHead.from_pretrained(model)
        # 对 value head 模型进行额外补丁处理
        patch_valuehead_model(model)

        if model_args.adapter_name_or_path is not None:
            vhead_path = model_args.adapter_name_or_path[-1]
        else:
            vhead_path = model_args.model_name_or_path

        # 加载 value head 的参数
        vhead_params = load_valuehead_params(vhead_path, model_args)
        
        if vhead_params is not None:
            # 加载 value head 参数到模型中
            model.load_state_dict(vhead_params, strict=False)
            logger.info_rank0(f"Loaded valuehead from checkpoint: {vhead_path}")

    # 如果模型不可训练，则冻结所有参数并切换为评估模式
    if not is_trainable:
        model.requires_grad_(False)
        for param in model.parameters():
            # 如果参数是 float32 但计算 dtype 不是，则转换类型
            if param.data.dtype == torch.float32 and model_args.compute_dtype != torch.float32:
                param.data = param.data.to(model_args.compute_dtype)

        model.eval()
    else:
        model.train()

    # 统计模型参数数量，用于日志输出
    trainable_params, all_param = count_parameters(model)
    if is_trainable:
        # 计算并格式化可训练参数比例
        param_stats = (
            f"trainable params: {trainable_params:,} || "
            f"all params: {all_param:,} || trainable%: {100 * trainable_params / all_param:.4f}"
        )
    else:
        # 仅显示总参数数
        param_stats = f"all params: {all_param:,}"

    # 输出参数统计信息
    logger.info_rank0(param_stats)

    # 如果开启打印参数状态，则在 rank 0 上输出每个参数的详细信息
    if model_args.print_param_status and int(os.getenv("LOCAL_RANK", "0")) == 0:
        for name, param in model.named_parameters():
            print(f"name: {name}, dtype: {param.dtype}, device: {param.device}, trainable: {param.requires_grad}")

    # 返回最终构建好的模型对象
    return model
